// --------------------------------------------------------------------------------------------------------------------
// <copyright file="CHOMPExampleAgent.cs" company="Microsoft Corporation">
// The MIT License (MIT)
// 
// Copyright (c) 2014, Microsoft Corporation
// 
// Permission is hereby granted, free of charge, to any person obtaining a copy
//  of this software and associated documentation files (the "Software"), to deal
//  in the Software without restriction, including without limitation the rights
//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//  copies of the Software, and to permit persons to whom the Software is
//  furnished to do so, subject to the following conditions:
// 
// The above copyright notice and this permission notice shall be included in
//  all copies or substantial portions of the Software.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
//  THE SOFTWARE.
// </copyright>
// --------------------------------------------------------------------------------------------------------------------

namespace Test.Robotics.Manipulation
{
    using System;
    using System.Collections.Generic;
    using System.Runtime.Serialization;
    using System.Threading;
    using System.Threading.Tasks;
    using System.Windows.Media.Media3D;

    using Microsoft.Robotics.Geometry;
    using Microsoft.Robotics.Geometry.SDF;
    using Microsoft.Robotics.Hardware.KukaLBR;
    using Microsoft.Robotics.Hardware.Mico;
    using Microsoft.Robotics.Kinematics;
    using Microsoft.Robotics.Manipulation;
    using Microsoft.Robotics.Manipulation.MotionPlanning;
    using Microsoft.Robotics.Manipulation.Runtime;
    using Microsoft.Robotics.Numerics;
    using Microsoft.Robotics.Runtime;
    using Microsoft.Robotics.Visualization;
    using ShoNS.Array;

    /// <summary>
    /// Test agent for the CHOMP motion planner. It takes a visualizer as input, 
    /// plans a feasible motion for a specific environment and sends the trajectory to the visualizer for display.
    /// The agent creates a set of parameters for CHOMP (<see cref="CHOMPparameters" /> class), creates a test environment (including obstacles, start and goal, initial path etc),
    /// generates a delegate function that computes the obstacle cost and gradient and passes all this to the constructor (MotionPlanCHOMP.cs) to create the motion planner. 
    /// It then starts up the visualizer and runs the planning algorithm until it converges (multiple options available in <see cref="CHOMPparameters" /> class). 
    /// Finally, it passes the output trajectory from the planner to the visualizer for display.
    /// Currently, 3 different test cases have been implemented:
    /// 1) An environment with a table at a specified height. The cost and gradient are computed in closed form. The robot's initial trajectory passes through the table.
    /// 2) An environment with 2 spherical obstacles. Again, cost and gradient are computed in closed form assuming known sphere centers and radii. Initial trajectory passes through the sphere.
    /// 3) An environment represented as a Signed Distance Function (generated from a mesh file). Cost is computed based on the signed distance and gradient by finite differencing. Initial trajectory passes through the objects in the environment.
    /// </summary>
    [DataContract]
    public class CHOMPExampleAgent : Agent, IDisposable
    {
        /// <summary>
        /// Instance of the CHOMP planner
        /// </summary>
        private MotionPlanCHOMP planner;

        /// <summary>
        /// Robot Kinematic chain
        /// </summary>
        private KinematicChain kc;

        /// <summary>
        /// Name of the arm visualizer agent
        /// </summary>
        [DataMember]
        private string armVisualizerAgentName;

        /// <summary>
        /// Name of the environment visualizer agent
        /// </summary>
        [DataMember]
        private string envVisualizerAgentName;

        /// <summary>
        /// Name of the environment display agent
        /// </summary>
        [DataMember]
        private string envDisplayAgentName;

        /// <summary>
        /// Model of the arm to use for testing
        /// </summary>
        [DataMember]
        private ArmModel armModel;

        /// <summary>
        /// The arm visualizer agent
        /// </summary>
        private IAgentHost armVisualizerAgent;

        /// <summary>
        /// The environment visualizer agent
        /// </summary>
        private IAgentHost envVisualizerAgent;

        /// <summary>
        /// The actual low level agent that displays the environment
        /// </summary>
        private IAgentHost envDisplayAgent;

        /// <summary>
        /// Cancellation token source used to cancel reading task.
        /// </summary>
        private CancellationTokenSource cts;

        /// <summary>
        /// Task for processing feedback.
        /// </summary>
        private Task planTask;

        /// <summary>
        /// Initial trajectory used to initialize CHOMP
        /// </summary>
        private DoubleArray initTrajectory;

        /// <summary>
        /// Parameter set used for CHOMP
        /// </summary>
        private CHOMPparameters initParams;

        /// <summary>
        /// Environment representation for CHOMP
        /// </summary>
        private CHOMPEnvironmentRepresentation environment;

        /// <summary>
        /// Initializes a new instance of the <see cref="CHOMPExampleAgent" /> class
        /// Creates a simple test for CHOMP and sends the output trajectory to the visualizer as a  SIMULATEDEETRAJECTORYMESSAGE and SIMULATEDJOINTTRAJECTORYMESSAGE
        /// </summary>
        /// <param name="name"> Name of the CHOMP Test Agent</param>
        /// <param name="armVisualizerAgentName"> Name of the arm Visualizer Agent </param>
        /// <param name="envVisualizerAgentName"> Name of the environment visualizer agent</param>
        /// <param name="envDisplayAgentName"> Name of the actual display agent for the environment </param>
        /// <param name="arm"> Arm model to be used for planning </param>
        public CHOMPExampleAgent(string name, string armVisualizerAgentName, string envVisualizerAgentName, string envDisplayAgentName, ArmModel arm)
            : base(name)
        {
            // Visualizer agent name
            this.armVisualizerAgentName = armVisualizerAgentName;
            this.envVisualizerAgentName = envVisualizerAgentName;
            this.envDisplayAgentName = envDisplayAgentName;
            this.armModel = arm;
        }

        /// <summary>
        /// Initializes the parameters for the planner, chooses an obstacle representation, cost and gradient delegate function,
        /// creates and starts up the planner.
        /// </summary>
        /// <param name="locator">Agent locator</param>
        public override void Initialize(AgentLocator locator)
        {
            base.Initialize(locator);

            if (this.armModel == ArmModel.KUKA)
            {
                // Create a KUKA kinemnatic chain
                this.kc = KukaLBR4KinematicConfigurations.CreateKuka7DOFKinematicChain();
            }
            else if (this.armModel == ArmModel.MICO)
            {
                // Create a Mico kinematic chain
                this.kc = MicoKinematicConfigurations.CreateMicoKinematicChain();
            }
            else
            {
                throw new ArgumentException("Unknown robot model :{0} chosen", this.armModel.ToString());
            }

            // Create an instance of the CHOMP parameters class with default values
            this.initParams = new CHOMPparameters();
            this.initParams.ObsWeight = 1.0;
            this.initParams.SmoothWeight = 1.0;
            this.initParams.StepSize = 1.0 / 250;
            this.initParams.NumPoints = 101; // keep this set
            this.initParams.SelfCollisionFactor = 20;
            this.initParams.EnvCollisionFactor = 200;
            this.initParams.EnvCollisionTolerance = 0.05; // this is passed in to the delegate function
            this.initParams.SelfCollisionTolerance = 0.02; // 2 cm
            this.initParams.UseSelfCollisionInConvergenceTesting = true;

            // Set convergence options
            this.initParams.ExitOnNoCollision = true;
            this.initParams.MaxIterations = 100;
            this.initParams.MinCostDifference = 0.00005;
            this.initParams.MaxTime = 10.0;

            // Test on scene with spherical obstacles
            ////this.SetupSceneWithSphericalObstacles();

            // OR TEST scene with an SDF from synthetic data
            this.SetupSceneWithSDF();

            // OR TEST scene with a single table obstacle and a cost function that restricts the arm to stay above a certain "z" value
            ////this.initParams.StepSize = 1.0 / 5000; // needs a small enough step size
            ////this.SetupSceneWithTable();

            // Initialize CHOMP - use function with sphere based cost and gradient as delegate
            this.planner = new MotionPlanCHOMP(
                this.kc,
                this.initParams,
                this.environment.ComputeSphereCostandGradient);

            // Get name of visualizer agent
            this.armVisualizerAgent = locator.GetAgent(this.armVisualizerAgentName);
            this.envVisualizerAgent = locator.GetAgent(this.envVisualizerAgentName);
            this.envDisplayAgent = locator.GetAgent(this.envDisplayAgentName);

            this.cts = new CancellationTokenSource();

            this.planTask = new Task((obj) => this.PlanPath(obj), this.cts.Token);
            this.planTask.Start();
        }

        /// <summary>
        /// Disposes of an instance of the <see cref="CHOMPExampleAgent"/> class.
        /// </summary>
        public void Dispose()
        {
            this.Dispose(true);
        }

        /// <summary>
        /// Disposes of an instance of the <see cref="CHOMPExampleAgent"/> class.
        /// </summary>
        /// <param name="disposing">True if disposing, false if finalizing.</param>
        private void Dispose(bool disposing)
        {
            if (disposing)
            {
                if (null != this.planner)
                {
                    this.planner.Dispose();
                    this.planner = null;
                }

                if (null != this.initTrajectory)
                {
                    this.initTrajectory.Dispose();
                    this.initTrajectory = null;
                }

                if (null != this.cts)
                {
                    this.cts.Dispose();
                    this.cts = null;
                }

                if (null != this.planTask)
                {
                    this.planTask.Dispose();
                    this.planTask = null;
                }
            }
        }

        /// <summary>
        /// Setup Test scene with Spherical obstacles. Cost is proportional to signed distance from sphere surfaces, gradients in closed form.
        /// </summary>
        private void SetupSceneWithSphericalObstacles()
        {
            // Generate spherical obstacles - TWO SPHERES COLLIDING WITH THE INITIAL TRAJECTORY
            List<GeometryElement3D> objects = new List<GeometryElement3D>();

            Vector3 location1 = new Vector3(-0.05, 0, 1.1);
            Vector3 location2 = new Vector3(-0.35, 0.1, 0.8);

            GeometryElement3D obstacle1 = new SphereElement(location1, 0.2);
            objects.Add(obstacle1);

            GeometryElement3D obstacle2 = new SphereElement(location2, 0.2);
            objects.Add(obstacle2);

            // Setup the environment
            this.environment = new SphereEnvironment();
            SphereEnvironment sphEnv = (SphereEnvironment)this.environment;
            sphEnv.Obstacles = new GeometryElementGroup3D(objects);

            // GENERATE A TRAJECTORY THAT COLLIDES WITH THE OBSTACLES
            DoubleArray targetJoint, currentJoint;
            if (this.armModel == ArmModel.KUKA)
            {
                targetJoint = DoubleArray.From(new double[] { 0, MathConstants.PIOverTwo, 0, 0, 0, 0, 0 }).Transpose();
                currentJoint = DoubleArray.From(new double[] { 0, -MathConstants.PIOverTwo, 0, 0, 0, 0, 0 }).Transpose();
            }
            else if (this.armModel == ArmModel.MICO)
            {
                // Empirically found out that step size needs to be large for this case.
                // This may be generic to the MICO, but further checks are needed to establish that
                this.initParams.StepSize = 1.0 / 50;

                // Mico
                targetJoint = DoubleArray.From(new double[] { 0, MathConstants.PIOverTwo, Math.PI, 0, 0, 0 }).Transpose();
                currentJoint = DoubleArray.From(new double[] { 0, Math.PI + MathConstants.PIOverTwo, Math.PI, 0, 0, 0 }).Transpose();
            }
            else
            {
                throw new ArgumentException("Unknown robot model : {0} input", this.armModel.ToString());
            }

            this.initTrajectory = MotionPlanCHOMP.GenerateTrajectoryUsingLinearInterpolation(currentJoint, targetJoint, this.initParams.NumPoints);
        }

        /// <summary>
        /// Setup Test scene with an SDF. Cost function is based on signed distance from SDF. Gradients are computed directly in the SDF
        /// </summary>
        private void SetupSceneWithSDF()
        {
            // Setup the environment and create the SDF
            this.environment = new SDFEnvironment();
            SDFEnvironment sdfEnv = (SDFEnvironment)this.environment;
            sdfEnv.Sdf = SDFIO.LoadSDF(@"Data\", "TableTop3");

            // GENERATE A TRAJECTORY THAT COLLIDES WITH THE OBSTACLES
            DoubleArray targetJoint, currentJoint;
            if (this.armModel == ArmModel.KUKA)
            {
                // KUKA TEST FOR SDF WITH MONKEY
                currentJoint = DoubleArray.From(new double[] { -Math.PI * 0.75, -Math.PI / 6, -Math.PI / 4, 0, 0, 0, 0 }).Transpose();
                targetJoint = DoubleArray.From(new double[] { Math.PI * 2 / 3, -Math.PI / 3, Math.PI / 2, 0, 0, 0, 0 }).Transpose();
            }
            else if (this.armModel == ArmModel.MICO)
            {
                // MICO TEST FOR SDF WITH MONKEY - TableTop3, both targets work
                currentJoint = DoubleArray.From(new double[] { Math.PI - Math.PI / 4, Math.PI / 2, Math.PI, 0, 0, 0 }).Transpose();
                targetJoint = DoubleArray.From(new double[] { Math.PI / 4, Math.PI, Math.PI / 2, 0, 0, 0 }).Transpose();
            }
            else
            {
                throw new ArgumentException("Unknown robot model : {0} input", this.armModel.ToString());
            }

            this.initTrajectory = MotionPlanCHOMP.GenerateTrajectoryUsingLinearInterpolation(currentJoint, targetJoint, this.initParams.NumPoints);
        }

        /// <summary>
        /// Setup Test scene with a single table obstacle. Cost and gradient are based on distance to table surface. Robot is restricted from going below table.
        /// </summary>
        private void SetupSceneWithTable()
        {
            //  CREATE AN ENVIRONMENT WITH A TABLE AT Z = 0 
            this.environment = new TableEnvironment();
            TableEnvironment tableEnv = (TableEnvironment)this.environment;
            tableEnv.TableHeight = 0.0; // table is at z = 0

            // CREATE AN ENVIRONMENT WITH A TABLE AT Z = 0 and an initial trajectory that collides with the table
            DoubleArray start, midPoint, goal;
            if (this.armModel == ArmModel.KUKA)
            {
                start = DoubleArray.From(new double[] { 0, MathConstants.PIOverTwo, 0, 0, 0, 0, 0 }).Transpose();
                midPoint = DoubleArray.From(new double[] { MathConstants.PIOverFour, MathConstants.PIOverTwo + MathConstants.PIOverFour, 0, -MathConstants.TwoPI / 12, 0, 0, 0 }).Transpose();
                goal = DoubleArray.From(new double[] { MathConstants.PIOverTwo, MathConstants.PIOverTwo, 0, 0, 0, 0, 0 }).Transpose();
            }
            else if (this.armModel == ArmModel.MICO) 
            {
                // VALUES NOT GOOD FOR MICO
                start = DoubleArray.From(new double[] { 0, MathConstants.PIOverTwo, Math.PI, 0, 0, 0 }).Transpose();
                midPoint = DoubleArray.From(new double[] { MathConstants.PIOverFour, MathConstants.PIOverTwo - MathConstants.PIOverFour, Math.PI, -MathConstants.TwoPI / 12, 0, 0 }).Transpose();
                goal = DoubleArray.From(new double[] { MathConstants.PIOverTwo, MathConstants.PIOverTwo, Math.PI, 0, 0, 0 }).Transpose();
            }
            else
            {
                throw new ArgumentException("Unknown robot model : {0} input", this.armModel.ToString());
            }

            // Linearly interpolate
            int pt1Ct, pt2Ct;
            if (this.initParams.NumPoints % 2 == 0)
            {
                pt1Ct = this.initParams.NumPoints / 2 + 1;
                pt2Ct = this.initParams.NumPoints / 2;
            }
            else
            {
                pt1Ct = pt2Ct = (int)(this.initParams.NumPoints / 2) + 1;
            }

            DoubleArray pt1 = MotionPlanCHOMP.GenerateTrajectoryUsingLinearInterpolation(start, midPoint, pt1Ct);
            DoubleArray pt2 = MotionPlanCHOMP.GenerateTrajectoryUsingLinearInterpolation(midPoint, goal, pt2Ct);

            this.initTrajectory = new DoubleArray(pt1.size0, pt1.size1 + pt2.size1 - 1);
            for (int j = 0; j < pt1.size0; j++)
            {
                for (int i = 0; i < pt1.size1; i++)
                {
                    this.initTrajectory[j, i] = pt1[j, i];
                }

                for (int i = 1; i < pt2.size1; i++)
                {
                    this.initTrajectory[j, (i - 1) + pt1.size1] = pt2[j, i];
                }
            }
        }

        /// <summary>
        /// Runs the motion planner using the "task" framework. After planning, it sends the initial and final trajectories to the visualizer for display.
        /// The function keeps replaying the final trajectory  and never exits.
        /// Parameter should be a <c>CancellationToken</c>.
        /// </summary>
        /// <param name="obj">A cancellation token.</param>
        private void PlanPath(object obj)
        {
            // Run CHOMP
            System.Console.WriteLine("Planning a collision-free path with CHOMP");
            DoubleArray finalTrajectory = this.planner.RunCHOMPOptimization(this.initTrajectory);
            System.Console.WriteLine("Planning done. Passing trajectory to visualizer");

            // In case we have obstacles initialized, try to visualize them
            if (this.environment is SphereEnvironment)
            {
                VisualizeGeometryElementGroup3DMessage msg = new VisualizeGeometryElementGroup3DMessage(((SphereEnvironment)this.environment).Obstacles);
                this.envVisualizerAgent.GetControlPort<VisualizeGeometryElementGroup3DMessage>().Post(msg);
            }

            // In case we have an SDF initialized, get the surface points for visualization
            if (this.environment is SDFEnvironment)
            {
                // Get the points on the SDF
                List<Tuple<Vector3, Vector3>> surface = SDFIO.GetSurfaceFromSDF(((SDFEnvironment)this.environment).Sdf, 0.5 * ((SDFEnvironment)this.environment).Sdf.Resolution);
                List<Point3D> surfacePoints3D = new List<Point3D>();
                for (int i = 0; i < surface.Count; i += 5)
                {
                    surfacePoints3D.Add(new Point3D(surface[i].Item1.X, surface[i].Item1.Y, surface[i].Item1.Z));
                }

                System.Console.WriteLine("Num points: {0}, Subsampled: {1}", surface.Count, surfacePoints3D.Count);

                // Send the points directly to the display agent for visualization
                UpdatePointCloudMessage msg = new UpdatePointCloudMessage(surfacePoints3D);
                this.envDisplayAgent.GetControlPort<UpdatePointCloudMessage>().Post(msg);
            }
 
            // Get trajectory waypoints
            List<double[]> initWayPoints = MotionPlanCHOMP.GetWayPointList(this.initTrajectory);

            // Send the initial waypoints to the visualizer
            SimulateEETrajectoryMessage initUpdateMsgFrames = new SimulateEETrajectoryMessage(initWayPoints);
            this.armVisualizerAgent.GetControlPort<SimulateEETrajectoryMessage>().Post(initUpdateMsgFrames);

            SimulateJointTrajectoryMessage initUpdateMsg = new SimulateJointTrajectoryMessage(initWayPoints);
            this.armVisualizerAgent.GetControlPort<SimulateJointTrajectoryMessage>().Post(initUpdateMsg);

            // Sleep for a few seconds to wait till the initial trajectory has finished displaying
            Thread.Sleep(7000);

            // Get trajectory waypoints
            List<double[]> finalWayPoints = MotionPlanCHOMP.GetWayPointList(finalTrajectory.T, false);

            while (true)
            {
                // Send the final waypoints to the visualizer
                SimulateEETrajectoryMessage finalUpdateMsgFrames = new SimulateEETrajectoryMessage(finalWayPoints);
                this.armVisualizerAgent.GetControlPort<SimulateEETrajectoryMessage>().Post(finalUpdateMsgFrames);

                SimulateJointTrajectoryMessage finalUpdateMsg = new SimulateJointTrajectoryMessage(finalWayPoints);
                this.armVisualizerAgent.GetControlPort<SimulateJointTrajectoryMessage>().Post(finalUpdateMsg);

                System.Console.ReadLine();
            }
        }
    }
}
